Efficient analytic gradient of the LMM objective + gradient-based optimizers - #903

Summary

This PR adds an efficient analytic gradient of the profiled objective (deviance / REML criterion) for LinearMixedModel, exposes it as the exported function objective_gradient!, and wires it into the NLopt backend so that the gradient-based optimizers :LD_LBFGS, :LD_MMA, and :LD_SLSQP can be selected via fit(MixedModel, form, data; optimizer=:LD_LBFGS). The default optimizer remains the derivative-free :LN_NEWUOA.

It also reworks the ForwardDiff extension to reuse the core linear-algebra pipeline instead of maintaining parallel fd_* implementations, and adds a gradient keyword to fit/fit! (:analytic default, :forwarddiff opt-in) so that ForwardDiff can serve as an alternative gradient source for the LD optimizers — see the sections below the BLAS-3 one.

Approach

The objective is an affine function of the logarithms of the diagonal elements of the blocked Cholesky factor, obj = 2 Σⱼ wⱼ log Lⱼⱼ + const, with weights 1 on the random-effects rows, 1 (REML) or 0 (ML) on the fixed-effects rows, and n, n − p, or pwrss/σ² on the ℓ_yy element for ML, REML, and fixed-σ objectives respectively. Differentiating Ω = LLᵀ (Murray 2016, arXiv:1602.07527) gives

∂obj/∂θₚ = tr(W L⁻¹ Ω̇ₚ L⁻ᵀ) = ⟨S, Ω̇ₚ⟩, S = L⁻ᵀWL⁻¹,

where S is shared by all parameters. A single blocked computation of the lower blocks of L⁻¹ (stored in a reusable GradientWorkspace), followed by weighted Gram products contracted against the sparse structure of the A blocks, therefore yields the entire gradient: Ω̇ₚ is never formed and no solve is repeated per parameter. Blocks between scalar terms with sparse A, and Diagonal/UniformBlockDiagonal diagonal blocks, use selected-entry kernels that evaluate only the entries of S that the contraction touches.

This replaces the earlier per-parameter draft in src/gradient.jl, which performed two full blocked triangular solves per parameter and disagreed with ForwardDiff for vector-valued terms (Λ where Λᵀ was needed in Omega_dot_diag_block!, and a missing Λᵣᵀ premultiplication of the off-diagonal blocks).

API

• objective_gradient!(g, m) — gradient at the current parameter values; returns the objective value.
• objective_gradient!(g, m, θ) — installs θ via setθ!/updateL! first, mirroring objective!.
• MixedModels.GradientWorkspace(m) (internal) — preallocated storage for allocation-free repeated evaluation inside the optimizer loop.
• Profiling a model fitted with an LD optimizer falls back to a derivative-free optimizer, since the profiling objectives do not evaluate gradients.
• The LD optimizers internally work on a per-observation scaling of the objective and gradient: on the deviance scale the gradient norm grows with n, making the identity a poor initial inverse-Hessian estimate, and the first LBFGS line-search step overshoots wildly for large data sets. fitlog, the progress display, and the reported fmin stay on the deviance scale; ftol_rel is invariant under the scaling and ftol_abs acts as a per-observation tolerance.

Performance

Gradient evaluation vs the ForwardDiff extension: ~60× faster on the maximal kb07 model, ~9× on insteval, with far less allocation.

Fits, gradient-based vs derivative-free:

Gradient-based optimization wins decisively on many-parameter models: mrk17 fits ~5× faster and reaches an objective 0.031 lower; kb07 fits ~7× faster. On ml1m (n = 10⁶, 2 parameters) LD_LBFGS needs only 13 evaluations but each gradient costs ~4× an objective evaluation, so wall time is on par with the derivative-free default; LD_MMA and LD_SLSQP also converge cleanly there (19 and 25 evaluations). Without the per-observation scaling, LBFGS failed outright on ml1m — see the second commit.

Dense-crossed cross term (BLAS-3)

For two scalar random-effects terms whose Cholesky factor has dense fill-in (the classic crossed-effects case, e.g. ml1m), the dominant cost of the gradient is the cross term S[r,b] on the sparsity pattern of A[r,b], previously one BLAS-1 column dot product per nonzero — memory-bandwidth bound (~2 Gflop/s measured). When the crossing is dense enough, this is now evaluated with a BLAS-3 matrix product X[r,r]ᵀX[r,b] a column-panel at a time (transient product bounded to q_r × 128, never fully materialized), running at ~80 Gflop/s. The path is gated on crossing density (nnz > 0.03·q_r·q_b) and dense fill, so sparse crossings (e.g. insteval) keep the BLAS-1 path; results are identical up to floating-point reassociation. On ml1m the gradient is ~1.24× faster.

Note on the original selected-inverse follow-up: a Takahashi-style selected inversion does not help these models. With dense Cholesky fill the factor pattern is dense, so forming X[r,b] directly is already cheaper than the selected-inverse alternative (which needs the full dense Z[r,r]), and the dominant selected-entry extraction cost is unchanged. The measurement that establishes this is in the commit message; the achievable win was the BLAS-3 restructuring above.

ForwardDiff extension rework

The extension previously duplicated the core objective pipeline in out-of-place fd_* variants (fd_setθ!, fd_updateL!, fd_pwrss, fd_logdet, fd_cholUnblocked!, fd_rankUpdate!). Every divergence from core was one of: array arguments instead of model fields, a BLAS/LAPACK call requiring BlasFloat, or a ::Float64/::T assertion blocking dual numbers. All six are now deleted in favor of:

• generic (non-BLAS) fallback methods for rankUpdate! (dense syrk! and UBD + BlockedSparse syr! paths; the BLAS methods are now T<:BlasFloat specializations and the dense-C signatures are narrowed to HermOrSym{T,Matrix{T}} to avoid ambiguity with the Diagonal/UniformBlockDiagonal-storage methods) and for cholUnblocked! (generic cholesky!(Hermitian(A, :L)); the Diagonal bound is relaxed from AbstractFloat to Real);
• array-based methods of setθ!, updateL!, pwrss, and a _logdet helper, with the model-based methods as thin wrappers.

Float64 dispatch still reaches the BLAS/LAPACK methods (verified with which); updateL! time and allocations on kb07 are identical to main, and no new method ambiguities are introduced.

Semantics change: fd_deviance now profiles σ (or holds it at optsum.sigma when fixed) and includes the constant weights term, exactly matching objective. ForwardDiff.gradient/hessian therefore refer to the profiled objective and agree with objective_gradient! at any θ (≤1e-12 on sleepstudy) without syncing the model state first — previously agreement relied on the envelope theorem at synced state. The Hessian reference values in test/forwarddiff.jl changed accordingly.

Gradient source selection

fit(MixedModel, form, data; optimizer=:LD_LBFGS, gradient=:forwarddiff) selects forward-mode AD as the gradient source; the default gradient=:analytic uses objective_gradient!. The option is stored in OptSummary (serialization-aware) and validated in fit!; requesting :forwarddiff without ForwardDiff.jl loaded throws an informative error. The ForwardDiff path caches the dual-promoted copies of A/L/reterms together with a GradientConfig and a DiffResults buffer in a reusable workspace (fd_gradient_workspace), so repeated gradient evaluations do not re-promote; value and gradient come from a single dual sweep. Both sources use the same per-observation scaling and take identical optimization paths (same evaluation counts, same optima).

Analytic vs ForwardDiff gradient source

From gradients/fd_vs_analytic.jl (one subprocess per configuration so peak RSS is per-configuration; measured on a refit! after a warm-up fit, so time/bytes/allocs exclude compilation):

The analytic gradient dominates on both axes: 55× faster than ForwardDiff on kb07 and 6× on ml1m, and — contrary to the intuition that the one-shot GradientWorkspace is the memory-heavy option — ForwardDiff also allocates more (1.9× the bytes on ml1m) and has a ~310 MiB higher peak RSS there, because the cached dual arrays carry nθ partials per entry. :forwarddiff is thus primarily a correctness oracle and fallback, not a memory-saving alternative.

Possible follow-ups

• Stream the dense fill block itself out of _invL! to bound peak memory (currently O(q_r·q_b) for the fill block); invasive to the hot path, memory-only benefit at scales beyond ml1m.
• Heuristics for when a gradient-based optimizer should be preferred (many covariance parameters, moderate n).
• Update profiling code to accept any optimizer

🤖 Generated with Claude Code